import os
import numpy as np
import cv2
import luojianet
import luojianet.nn as nn
import luojianet.ops as ops
from luojianet import nn, ops, Parameter, Tensor
from scipy.io import loadmat

def minibatch_sample(gt_mask, minibatch_size):
    rs = np.random.RandomState(5)

    cls_list = np.unique(gt_mask)
    cls_list = np.delete(cls_list, 0)
    # print(cls_list)
    idx_dict_per_class = dict()
    for cls in cls_list:

        train_idx_per_class = np.where(gt_mask == cls, gt_mask, np.zeros_like(gt_mask))
        idx = np.where(train_idx_per_class.ravel() == cls)[0]

        rs.shuffle(idx)

        idx_dict_per_class[cls] = idx

    train_idx_list = []
    count = 0
    while True:
        train_idx = np.zeros_like(gt_mask).ravel()
        for cls, idx in idx_dict_per_class.items():
            left = count * minibatch_size
            if left >= len(idx):
                continue
            # remain last batch though the real size is smaller than minibatch_size
            right = min((count + 1) * minibatch_size, len(idx))
            fetch_idx = idx[left:right]
            train_idx[fetch_idx] = cls
        count += 1
        if train_idx.sum() == 0:
            return train_idx_list
        train_idx_list.append(train_idx.reshape(gt_mask.shape))

def image_pad(image, encoder_size):
    image_shape = image.shape
    image_shape = np.array(image_shape)
    pad_shape = image_shape

    pad_shape[1] = int(np.ceil(image_shape[1] / encoder_size) * encoder_size)
    pad_shape[2] = int(np.ceil(image_shape[2] / encoder_size) * encoder_size)

    out = np.zeros((pad_shape[0], pad_shape[1], pad_shape[2]))
    out[:, 0:image.shape[1], 0:image.shape[2]] = image

    return out

def image_pad_(image, encoder_size):
    image_shape = image.shape
    image_shape = np.array(image_shape)
    pad_shape = image_shape

    pad_shape[0] = int(np.ceil(image_shape[0] / encoder_size) * encoder_size)
    pad_shape[1] = int(np.ceil(image_shape[1] / encoder_size) * encoder_size)

    out = np.zeros((pad_shape[0], pad_shape[1]))
    out[0:image.shape[0], 0:image.shape[1]] = image

    return out


class WhuHIDataset():
    num_classes = 9
    def __init__(self, mode, logger_handle, dataset_cfg):
        config = dataset_cfg
        if mode == 'TRAIN':
            self.gt = loadmat(config['train_gt_dir'])[config['train_gt_name']]
            self.data = loadmat(config['train_data_dir'])[config['train_data_name']]
        else:
            self.gt = loadmat(config['test_gt_dir'])[config['test_gt_name']]
            self.data = loadmat(config['test_data_dir'])[config['test_data_name']]
        im_cmean = self.data.reshape((-1, self.data.shape[-1])).mean(axis=0)
        im_cstd = self.data.reshape((-1, self.data.shape[-1])).std(axis=0)
        self.data = (self.data - im_cmean) / im_cstd

        self.data = self.data.transpose(2, 0, 1)

        self.data = image_pad(self.data, config['encoder_size'])
        # self.data = self.data.reshape(1, self.data.shape[0], self.data.shape[1], self.data.shape[2])
        self.data = Tensor.from_numpy(self.data.astype(np.float32))

        # self.gt = self.gt.reshape(1, self.gt.shape[0], self.gt.shape[1])
        self.gt = image_pad_(self.gt, config['encoder_size'])

        self.gt_idx = minibatch_sample(self.gt, 5)

    def __getitem__(self, index):
        temp = self.gt_idx[index].astype(np.float32) - 1

        return self.data, Tensor.from_numpy(temp)


    def __len__(self):
        return len(self.gt_idx)